#!/usr/bin/env python2.5
#coding: utf8

#Author: Gang Li
#Date: 2010-04-10 18:20

import socket
import sys
import select
import time
import cPickle as _p
from Queue import Queue as _q
from threading import Thread as _t
from threading import RLock
import logging

from core.utils import *
_r,  _w = Transporter.read_cmd, Transporter.write_cmd

from core.fs import FSRoot
from core.task import Task, SubTask

class Server:
    """
    free
    alloc
    wait_task
    """

    def __init__(self, binding_ip, port):
        self.port=port
        self.sock=socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.bind((binding_ip, port))
        #tcp use this line 
        self.sock.listen(1)

        self.nodes = {}

        self.last_check = 0

        self.waits = [self.sock]

        self.cid = 0
        self.cid_lock = RLock()
        self.task_lock = RLock()

        self.tasks = {}

        #recevied task
        self.waiting_tasks=[]

        self.fs = FSRoot.open(meta_sock = self.sock, mode="meta")
        self.fs.get_root().mkdirs("/x/y/z/u/w/i")

    def start_loop(self):
        print "waiting for connections ..."
        while 1:
            try:
                self.process()
                self.heart_beat()
                self.task_loop()
            except Exception,msg:
                print msg

    def allocate_free_nodes(self, size, gid):
        al = []

        if size > len(self.nodes):
            return

        self.task_lock.acquire()
        try:
            for n in self.nodes:
                node_state = self.nodes[n]
                if node_state.get("state") or node_state.get("state") == 'free':
                    al.append(n)

                    if len(al) == size:
                        for a in al:
                            self.nodes[a]['state'] = 'alloc'
                            logging.info("send alloc node cmd for %s"%gid)
                            _w(a, TASK_NODE_ALLOC, gid)
                        return al
        finally:
            self.task_lock.release()

    def task_loop_r(self, task):
        task.reduces = []
        for i in range(task.r):
            st = SubTask(task.id, task.name, type="reduce", res = task.res)
            st.id = hash(st)

            _w(task.nodes[task.m+i], START_SUB_TASK, st)
            task.reduces.append(st.id)
            print "start reduce func"
            yield st

    def task_loop(self):
        for task in self.waiting_tasks:
            nodes = self.allocate_free_nodes(task.m+task.r, task.id)

            if nodes:
                path = task.dir

                task.nodes = nodes
                self.tasks[task.id] = task

                #sub map task list
                task.maps=[]
                for i in range(task.m):
                    st = SubTask(task.id, task.name, type="map", res = task.res)
                    st.id = hash(st)

                    _w(nodes[i], START_SUB_TASK, st)
                    task.maps.append(st.id)

                task.start_reduce_task = self.task_loop_r(task)
                self.waiting_tasks.remove(task)

            else:
                logging.info("allocate task resources error, need:%s current: %s"%(task.m+task.r, len(self.nodes)))

    def process(self):
        n = select.select( self.waits, [], [], 2 )
        for skt in n[0]:
            if skt == self.sock:
                new_sock, (remote_ip, remote_port)=self.sock.accept()
                i,p = new_sock.getsockname()
                print "sockname",i,p
                logging.debug("new connection from %s:%s"%(remote_ip, remote_port))
                self.nodes[new_sock] = {}
                self.waits.append(new_sock)
            else:
                try:
                    if self.nodes[skt].get("initialized"):
                        id,len,data = _r(skt)

                        if id.startswith(FS_CMD):
                            if self.fs: self.fs.on_cmd(id, len, data)
                        elif id == HEART_BEAT_REQ:
                            _w(skt,HEART_BEAT_RSP)
                        elif id == HEART_BEAT_RSP:
                            pass
                        elif id == START_TASK:
                            #user data
                            print "start task"
                            task.id = hash(data)
                            self.waiting_tasks.append(data)
                            print "start task ok"
                        elif id == TASK_NODE_ALLOC_RSP:
                            self.nodes[skt]['state'] = 'wait_task'
                            gid=int(data)
                            logging.info("recv alloc node cmd rsp for %s"%gid)
                        elif id == START_SUB_TASK_OK:
                            self.nodes[skt]['state'] = 'running_task'
                        elif id == SUB_TASK_DONE:
                            self.nodes[skt]['state'] = 'free'
                            st = data

                            if hasattr(st, "error"):
                                print skt.getsockname(), st.error

                            task = self.tasks[st.pid]
                            
                            if st.type == 'map':
                                task.maps.remove(st.id)

                                if not task.maps:
                                    print "map done!!!"

                                for st in task.start_reduce_task:
                                    print "start reduce!!!"
                            elif st.type == 'reduce':
                                task.reduces.remove(st.id)
                                
                                if not task.reduces:
                                    print "reduce done!!!, task done."        
                                    task.nodes = []

                        self.nodes[skt]["last_seen"] = time.time()
                    else:
                        #initialize
                        if self.verify_client(skt):
                            self.nodes[skt]["initialized"]=True
                            self.nodes[skt]["state"] = 'free'
                            self.nodes[skt]["last_seen"] = time.time()
                        else:
                            self.clear_client(skt)
                except Exception, msg:
                    print "error: ",msg
                    self.clear_client(skt)

    def heart_beat(self):
        cur = time.time()
        if cur - self.last_check > 15:
            for skt in self.waits:
                if skt != self.sock:
                    last_seen = self.nodes[skt]["last_seen"]

                    if cur - last_seen > 60:
                        self.clear_client(skt)
                        print "dead node",skt
                    elif cur - last_seen > 20:
                        _w(skt,HEART_BEAT_REQ)
                        print "sent hb to node", skt
                    else:
                        print "health node", skt

            self.last_check = cur

    def get_cid(self):
        self.cid_lock.acquire()
        try:
            self.cid += 1
            return self.cid
        finally:
            self.cid_lock.release()

    def write_data(self,skt,data):
        _w(skt,DATA,data)
        self.nodes[skt]["last_seen"] = time.time()

    def clear_client(self, skt):
        print "clean dead node"
        skt.close()
        self.waits.remove(skt)
        del self.nodes[skt]

    def verify_client(self, skt):
        id,len,data = _r(skt)
        if id == VERIFY_REQ:
            _w(skt, VERIFY_RSP, self.get_cid())
        else:
            print "error verify cmd", id, len, data, "expect", VERIFY_REQ
            return

        id,len,data = _r(skt)
        if id == VERIFY_OK:
            return True

class Client:
    def __init__(self, srv_ip, srv_port):
        try:
            self.sock = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
            self.sock.connect((srv_ip,srv_port))
            self.last_check = time.time()
            self.last_seen = self.last_check
            self.r_queue = _q()
            self.w_queue = _q()
            
            self.socks = [self.sock]

            _t(target = self._start_loop, args = [self]).start()
            _t(target = self._start_dispatch, args= [self]).start()
        except socket.error, e:
            print "socket error"
            raise e

    def verify(self):
        _w(self.sock, VERIFY_REQ)

        id,len,data = _r(self.sock)
        if id == VERIFY_RSP:
            _w(self.sock, VERIFY_OK)
            logging.info("handshake success")
            return True
        else:
            logging.info("handshake failed")
            return

    def write(self, id, data=None): 
        self.w_queue.put((id, data))

    def _write(self):
        if self.w_queue.empty():return
        try:
            _w(self.sock,*self.w_queue.get())
            self.last_seen = time.time()
        except BaseException, e:
            logging.warning("write cmd error: %s" % e)

    def _start_loop(self, kw):
        if not self.verify():
            self.clear_server()
            return

        self.connected()

        while 1:
            try:
                self._process()
                self.heart_beat()
                self._write()
            except BaseException,msg:
                print msg
                self._force_heart_beat()

    def _start_dispatch(self, kw):
        while 1:
            try:
                self.on_cmd(*self.r_queue.get())
            except BaseException, e:
                logging.warning("dispatch msg error: %s"%e)

    def connected(self):pass

    def on_cmd(self, id, data): pass

    def _force_heart_beat(self):
        #resend client ID
        try:
            _w(self.sock,HEART_BEAT_REQ)
        except:
            self.clear_server()

    def clear_server(self):
        print "server has gone!!!"
        self.sock.close()
        self.socks.remove(self.sock)
        #TODO retry to connect server

    def _process(self):
        r,w,o = select.select(self.socks, [], [], 0.5)
        if r:
            for sock in r:
                id,len,data = _r(sock)
                self.r_queue.put((id, data))
                self.last_seen = time.time()

    def heart_beat(self):
        cur = time.time()
        if cur - self.last_check > 20:
            if cur - self.last_seen > 60:
                self.clear_server()
            elif cur - self.last_seen > 20:
                _w(self.sock,HEART_BEAT_REQ)
            else:
                pass

            self.last_check = cur

    def close(self):
        self.clear_server()